import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import DropPath
from timm.models.layers import to_2tuple

from .block import conv_block
from torch.utils.checkpoint import checkpoint


class Scale(nn.Module):
    """
    Scale vector by element multiplications.
    """

    def __init__(self, dim, init_value=1.0, trainable=True):
        super().__init__()
        self.scale = nn.Parameter(init_value * torch.ones(dim), requires_grad=trainable)

    def forward(self, x):
        return x * self.scale


class StarReLU(nn.Module):
    """
    StarReLU: s * relu(x) ** 2 + b
    """

    def __init__(self, scale_value=1.0, bias_value=0.0, scale_learnable=True, bias_learnable=True, mode=None, inplace=False):
        super().__init__()
        self.inplace = inplace
        self.relu = nn.ReLU(inplace=inplace)
        self.scale = nn.Parameter(scale_value * torch.ones(1), requires_grad=scale_learnable)
        self.bias = nn.Parameter(bias_value * torch.ones(1), requires_grad=bias_learnable)

    def forward(self, x):
        return self.scale * self.relu(x) ** 2 + self.bias


class SepConv(nn.Module):
    r"""
    Inverted separable convolution from MobileNetV2: https://arxiv.org/abs/1801.04381.
    """

    def __init__(self, dim, expansion_ratio=2, act1_layer=StarReLU, act2_layer=nn.Identity, bias=False, kernel_size=7, padding=3,
                 **kwargs, ):
        super().__init__()
        med_channels = int(expansion_ratio * dim)
        self.pwconv1 = nn.Linear(dim, med_channels, bias=bias)
        self.act1 = act1_layer()
        self.dwconv = nn.Conv2d(
            med_channels, med_channels, kernel_size=kernel_size,
            padding=padding, groups=med_channels, bias=bias)  # depthwise conv
        self.act2 = act2_layer()
        self.pwconv2 = nn.Linear(med_channels, dim, bias=bias)

    def forward(self, x):
        x = self.pwconv1(x)
        x = self.act1(x)
        x = x.permute(0, 3, 1, 2)
        x = self.dwconv(x)
        x = x.permute(0, 2, 3, 1)
        x = self.act2(x)
        x = self.pwconv2(x)
        return x


class Mlp(nn.Module):
    """ MLP as used in MetaFormer models, eg Transformer, MLP-Mixer, PoolFormer, MetaFormer baslines and related networks.
    Mostly copied from timm.
    """

    def __init__(self, dim, mlp_ratio=4, out_features=None, act_layer=StarReLU, drop=0., bias=False, **kwargs):
        super().__init__()
        in_features = dim
        out_features = out_features or in_features
        hidden_features = int(mlp_ratio * in_features)
        drop_probs = to_2tuple(drop)

        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
        self.act = act_layer()
        self.drop1 = nn.Dropout(drop_probs[0])
        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
        self.drop2 = nn.Dropout(drop_probs[1])

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop1(x)
        x = self.fc2(x)
        x = self.drop2(x)
        return x


class UpConvBlock(nn.Module):
    """ MLP as used in MetaFormer models, eg Transformer, MLP-Mixer, PoolFormer, MetaFormer baslines and related networks.
    Mostly copied from timm.
    """

    def __init__(self, dim, mlp_ratio=4, k_size=3, scale=2, out_features=None, act_layer=StarReLU, bias=False, **kwargs):
        super().__init__()
        in_features = dim
        out_features = out_features or in_features
        hidden_features = int(mlp_ratio * in_features)

        self.norm = nn.LayerNorm(in_features)
        self.fc0 = nn.Conv2d(in_features, in_features, kernel_size=k_size, padding=k_size // 2, bias=bias)
        self.up = nn.Upsample(scale_factor=scale, mode='bilinear') if scale!=1 else nn.Identity()
        self.fc1 = nn.Conv2d(in_features, hidden_features, kernel_size=k_size, padding=k_size // 2, bias=bias)
        self.act = act_layer()
        self.fc2 = nn.Conv2d(hidden_features, out_features, kernel_size=1, bias=bias)

    def forward(self, x):
        x = self.norm(x)
        x = x.permute(0, 3, 1, 2)  # [B,C,H,W]
        x = self.fc0(x)
        x = self.up(x)
        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)
        return x


class MetaFormerBlock(nn.Module):
    """
    Implementation of one MetaFormer block.
    """

    def __init__(self, dim,
                 token_mixer=SepConv, mlp=Mlp,
                 norm_layer=nn.LayerNorm,
                 drop=0., drop_path=0.,
                 layer_scale_init_value=None, res_scale_init_value=None
                 ):
        super().__init__()

        self.norm1 = norm_layer(dim)
        self.token_mixer = token_mixer(dim=dim, drop=drop)
        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.layer_scale1 = Scale(dim=dim, init_value=layer_scale_init_value) \
            if layer_scale_init_value else nn.Identity()
        self.res_scale1 = Scale(dim=dim, init_value=res_scale_init_value) \
            if res_scale_init_value else nn.Identity()

        self.norm2 = norm_layer(dim)
        self.mlp = mlp(dim=dim, drop=drop)
        self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.layer_scale2 = Scale(dim=dim, init_value=layer_scale_init_value) \
            if layer_scale_init_value else nn.Identity()
        self.res_scale2 = Scale(dim=dim, init_value=res_scale_init_value) \
            if res_scale_init_value else nn.Identity()

    def forward(self, x):
        x = self.res_scale1(x) + \
            self.layer_scale1(
                self.drop_path1(
                    self.token_mixer(self.norm1(x))
                )
            )
        x = self.res_scale2(x) + \
            self.layer_scale2(
                self.drop_path2(
                    self.mlp(self.norm2(x))
                )
            )
        return x


class DownBlock(nn.Module):
    """
    Top to Down Block for HourGlass Block
    Consist of ConvNet Block and Pooling
    """

    def __init__(self, ch_in=64, ch_out=64, n_block=3):
        super().__init__()

        self.former_block = nn.Sequential(*[MetaFormerBlock(dim=ch_in, res_scale_init_value=1.0) for i in range(n_block)])
        self.down_sample = nn.Sequential(
            nn.Upsample(scale_factor=0.5, mode='bilinear'),
            nn.Conv2d(ch_in, ch_out, kernel_size=3, padding=1),
        )

    def forward(self, x):
        x = self.former_block(x)

        x_down = x.permute(0, 3, 1, 2)
        x_down = self.down_sample(x_down)
        x_down = x_down.permute(0, 2, 3, 1)

        return x_down, x


class UpBlock(nn.Module):
    """
    Bottom Up Block for HourGlass Block
    Consist of ConvNet Block and Upsampling Block
    """

    def __init__(self, ch_in=64, ch_out=64, n_block=3):
        super().__init__()

        self.former_block = nn.Sequential(*[MetaFormerBlock(dim=ch_out, res_scale_init_value=1.0) for i in range(n_block)])
        self.up_sample = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear'),
            nn.Conv2d(ch_in, ch_out, kernel_size=3, padding=1),
        )

    def forward(self, x, res):
        x = x.permute(0, 3, 1, 2)
        x = self.up_sample(x)
        x = x.permute(0, 2, 3, 1)

        return self.former_block(x + res)


class HourGlassBlock(nn.Module):
    """
    Hour Glass Block for SR Model
    """

    def __init__(self, dims=(128, 256, 512, 640), n_down=2, n_up=3, n_mid=2, n_tail=2, grad_checkpointing=True):
        super(HourGlassBlock, self).__init__()
        self.n_tail = n_tail
        self.grad_checkpointing = grad_checkpointing

        self.down_blocks = nn.ModuleList([DownBlock(dims[i], dims[i + 1], n_block=n_down) for i in range(len(dims) - 1)])
        self.mid_blocks = nn.Sequential(*[MetaFormerBlock(dim=dims[-1], res_scale_init_value=1.0) for i in range(n_mid)])
        self.skip_convs = nn.ModuleList([Mlp(dims[i], 2) for i in range(len(dims) - 1)])
        self.skip_convs.insert(0, Mlp(dims[0], 2))
        self.up_blocks = nn.ModuleList([UpBlock(dims[i], dims[i - 1], n_block=n_up) for i in range(len(dims) - 1, 0, -1)])

        if n_tail > 0:
            self.tail = nn.Sequential(*[MetaFormerBlock(dim=dims[0], res_scale_init_value=1.0) for i in range(n_tail)])

    def forward_ckpt(self, x):
        res_list = []
        out = x
        for down in self.down_blocks:
            out, res = checkpoint(down, out)
            res_list.append(res)
        
        out = checkpoint(self.mid_blocks, out)

        def up_forward(up, out, skip, res):
            return up(out, skip(res))

        for up, skip, res in zip(self.up_blocks, self.skip_convs[:0:-1], res_list[::-1]):
            out = checkpoint(up_forward, up, out, skip, res)

        out_inter = self.skip_convs[0](x) + out

        if self.n_tail > 0:
            out = checkpoint(self.tail, out_inter)
        else:
            out = out_inter
        # Change to Residul Structure
        return out, out_inter


    def forward(self, x):
        if self.grad_checkpointing and self.training:
            return self.forward_ckpt(x)

        res_list = []
        out = x
        for down in self.down_blocks:
            out, res = down(out)
            res_list.append(res)

        out = self.mid_blocks(out)

        for up, skip, res in zip(self.up_blocks, self.skip_convs[:0:-1], res_list[::-1]):
            out = up(out, skip(res))

        out_inter = self.skip_convs[0](x) + out

        if self.n_tail > 0:
            out = self.tail(out_inter)
        else:
            out = out_inter
        # Change to Residul Structure
        return out, out_inter


class HourGlassNetMultiScaleFormer(nn.Module):
    """
    Hour Glass SR Model, Use Mutil-Scale Label(HR_down_Xn) Supervision.
    """

    def __init__(self, in_nc=3, out_nc=3, upscale=4, dims=(96, 192, 320, 640), n_mid=2, n_tail=2, HG_stage=(1, 3, 5),
                 inter_supervis=True):
        super().__init__()

        self.HG_stage = HG_stage
        n_HG = HG_stage[-1]
        self.inter_supervis = inter_supervis

        self.conv_in = conv_block(in_nc, dims[0], kernel_size=3, norm_type=None, act_type=None)

        # Actually, the ksize can be any size for all scales
        self.flat_map = UpConvBlock(dim=dims[0], out_features=out_nc, scale=upscale)
        self.edge_map = UpConvBlock(dim=dims[0], out_features=out_nc, scale=upscale)
        self.corner_map = UpConvBlock(dim=dims[0], out_features=out_nc, scale=upscale)

        self.upsample_flat = UpConvBlock(dim=dims[0], out_features=out_nc, scale=upscale)
        self.upsample_edge = UpConvBlock(dim=dims[0], out_features=out_nc, scale=upscale)
        self.upsample_corner = UpConvBlock(dim=dims[0], out_features=out_nc, scale=upscale)

        self.HG_blocks = nn.ModuleList([HourGlassBlock(dims, n_mid=n_mid, n_tail=n_tail if i != n_HG else 0) for i in range(n_HG + 1)])

    def forward(self, x):
        x = self.conv_in(x)
        x = x.permute(0, 2, 3, 1)  # [B,H,W,C]

        SR_map = []
        result = []
        out = x

        # Multi-Scale supervise, [2, 2, 2] for 6, [2, 3, 3] for 8
        super_block_idx = self.HG_stage

        for i, HG in enumerate(self.HG_blocks):
            out, out_inter = HG(out)

            if i in super_block_idx:
                sr_feature = out.mul(0.2) + x

                if i == super_block_idx[0]:
                    srout_flat = self.upsample_flat(sr_feature)
                    flat_map = self.flat_map(sr_feature)
                    result.append(srout_flat)
                elif i == super_block_idx[1]:
                    srout_edge = self.upsample_edge(sr_feature)
                    edge_map = self.edge_map(sr_feature)
                    result.append(srout_edge)
                elif i == super_block_idx[2]:
                    srout_corner = self.upsample_corner(sr_feature)
                    corner_map = self.corner_map(sr_feature)
                    result.append(srout_corner)

                    maps = torch.stack((flat_map, edge_map, corner_map)) # [N,B,C,H,W]
                    maps = F.softmax(maps, dim=0)
                    out = maps[0] * srout_flat + maps[1] * srout_edge + maps[2] * srout_corner
                    result.append(out)

                    for map in maps:
                        SR_map.append(map.mean(dim=1, keepdim=True))


        return result, SR_map


if __name__ == '__main__':
    from torchanalyzer import ModelFlopsAnalyzer, TorchViser

    model = HourGlassNetMultiScaleFormer(in_nc=3, out_nc=3, upscale=2, HG_stage=(3, 4, 5), inter_supervis=True)

    inputs = torch.randn(1, 3, 64, 64)

    analyzer = ModelFlopsAnalyzer(model)
    info = analyzer.analyze(inputs)
    TorchViser().show(model, info)
